"""NVIDIA GPU 功耗采样后端（通过 NVML）"""

from __future__ import annotations
from typing import Optional

from .base import PowerBackend


# 导入 NVML 库（优先使用 nvidia-ml-py，向后兼容 pynvml）
def _import_nvml():
    """
    导入 NVML 库，优先使用 nvidia-ml-py（新包名），
    如果不存在则回退到 pynvml（向后兼容）
    
    Returns:
        导入的 NVML 模块
    """
    try:
        # 优先尝试 nvidia-ml-py（新包名，不会产生废弃警告）
        import nvidia_ml_py as pynvml
        return pynvml
    except ImportError:
        # 回退到 pynvml（旧包名，向后兼容）
        import pynvml
        return pynvml


class NVMLBackend(PowerBackend):
    """NVIDIA GPU 功耗采样（通过 NVML）"""
    
    def __init__(self):
        self._handle = None
        self._nvml = None  # 缓存的 NVML 模块
    
    @staticmethod
    def is_available() -> bool:
        try:
            pynvml = _import_nvml()
            pynvml.nvmlInit()
            pynvml.nvmlShutdown()
            return True
        except:
            return False
    
    def init(self, device_id: int = 0) -> bool:
        try:
            self._nvml = _import_nvml()
            self._nvml.nvmlInit()
            self._handle = self._nvml.nvmlDeviceGetHandleByIndex(device_id)
            return True
        except Exception as e:
            print(f"NVML init failed: {e}")
            return False
    
    def sample(self) -> Optional[float]:
        try:
            if self._nvml is None:
                self._nvml = _import_nvml()
            power_mW = self._nvml.nvmlDeviceGetPowerUsage(self._handle)
            return float(power_mW) if power_mW > 0 else None
        except:
            return None
    
    def cleanup(self):
        try:
            if self._nvml is not None:
                self._nvml.nvmlShutdown()
        except:
            pass
    
    def get_name(self) -> str:
        return "nvml"

